"""
用分支限界法解决旅行商问题。
作者：ipid （ https://github.com/ipid ）
"""

from dataclasses import dataclass
from queue import PriorityQueue
from typing import List

from tsp_common_code import read_graph, AdjacencyMatrix


@dataclass
class SearchState:
    route: List[int]
    visited: List[int]
    dist: int
    value: int

    def __init__(self, *, route: List[int], visited: List[int], dist: int):
        self.route = route
        self.visited = visited
        self.dist = dist
        self.value = -1

    def __lt__(self, other: "SearchState") -> int:
        return self.value < other.value


class BranchAndBoundSolution:
    graph: AdjacencyMatrix
    # 分支限界法的上界
    upper_limit: int

    # 最佳路线
    best_route: List[int]

    def __init__(self, graph: AdjacencyMatrix):
        self.graph = graph

    def greedy(self) -> None:
        """
        通过贪心算法求出上界
        :return: None
        """
        city_num = len(self.graph)
        visited = [False] * city_num
        visited[0] = True
        self.best_route = [0]

        # 执行 city_num - 1 次循环
        curr_city: int = 0
        total_dist = 0
        for i in range(city_num - 1):
            next_city = min((dist, index) for index, dist in enumerate(self.graph[curr_city]) if not visited[index])[1]
            self.best_route.append(next_city)
            visited[next_city] = True
            total_dist += self.graph[curr_city][next_city]
            curr_city = next_city

        self.best_route.append(0)
        total_dist += self.graph[curr_city][0]
        self.upper_limit = total_dist

    def calculate_value_of_state(self, state: SearchState) -> None:
        """
        通过状态计算估值函数的数值
        :param state: 状态
        :return: 估值
        """
        state.value = state.dist
        for i in range(len(self.graph)):
            if not state.visited[i]:
                state.value += min(dist for index, dist in enumerate(self.graph[i]) if index != i)

        state.value += min(self.graph[0])

    def solve(self) -> None:
        # 用贪心算法求出上界
        self.greedy()

        queue: "PriorityQueue[SearchState]" = PriorityQueue()
        initial_state = SearchState(
            route=[0],
            visited=[True] + [False] * (len(self.graph) - 1),
            dist=0,
        )
        self.calculate_value_of_state(initial_state)
        queue.put(initial_state)

        while not queue.empty():
            state = queue.get()
            print(state)

            if len(state.route) == len(self.graph):
                new_dist = state.dist + self.graph[state.route[-1]][0]
                new_route = state.route + [0]
                new_state = SearchState(
                    route=new_route,
                    visited=state.visited[:],
                    dist=new_dist
                )
                new_state.value = new_dist
                queue.put(new_state)
                continue

            if len(state.route) == len(self.graph) + 1:
                self.upper_limit = state.dist
                self.best_route = state.route
                return

            for i in range(len(self.graph)):
                if not state.visited[i]:
                    new_route = state.route + [i]
                    new_dist = state.dist + self.graph[state.route[-1]][i]
                    new_visited = state.visited[:]
                    new_visited[i] = True

                    new_state = SearchState(
                        route=new_route,
                        visited=new_visited,
                        dist=new_dist,
                    )
                    self.calculate_value_of_state(new_state)
                    if new_state.value < self.upper_limit:
                        queue.put(new_state)


def main():
    graph = read_graph()
    if len(graph) <= 1:
        result_dist, result_route = 0, []
    else:
        solution = BranchAndBoundSolution(graph)
        solution.solve()
        result_dist, result_route = solution.upper_limit, solution.best_route

    print(f'\n最佳路线：{result_route}\n最佳路线总长度：{result_dist}\n')


if __name__ == '__main__':
    main()
